from llama_index.core.prompts import PromptTemplate
from llama_index.llms.openai import OpenAI
from fastapi import APIRouter
import yaml

router = APIRouter(
    prefix="/gj-rag",
    tags=["query"],
    responses={404: {"description": "Not found"}}
)
#获取配置信息
with open('config.yaml', 'r') as file:
        config = yaml.load(file, Loader=yaml.FullLoader)

llm = OpenAI(model=config["openai"]["model"], api_base=config["openai"]["api_base"], api_key=config["openai"]["api_key"])


query_gen_str = """\
你是一个专业帮助用户优化查询的助手，请你将用户的查询改写或转述为用语清晰、通俗易懂、目的明确的简体现代汉语，务必确保改写后的查询用语标准规范，\
如果用户的查询需要进行一步以上的推理才能获得答案，请将用户的查询拆分以便于推理，例如：用户询问：“《易经全转》的作者来自哪里？”，这个问题应该拆分成“《易经全转》的作者是谁？《易经全转》的作者来自哪里？”。\
如果用户的提问明显包含两个以上的问题，请将用户的查询拆分成子问题，例如：用户询问：“人们如何评价《子夏易传》的内容和作者？”，这个问题应该拆分成“人们如何评价《子夏易传》的内容？《子夏易传》的作者是谁？人们如何评价《子夏易传》的作者？”。\
请直接返回你改写后的问题，不要编号，不要换行符"\n"，除了问题本身不要其他任何内容。\
用户查询如下: {query}
"""
query_gen_prompt = PromptTemplate(query_gen_str)

def generate_queries(query: str):
    response = llm.predict(
        query_gen_prompt, query=query
    )
    # assume LLM proper put each query on a newline
    # queries = response.split("\n")
    # queries_str = "\n".join(queries)
    # print(f"Generated queries:\n{response}")
    return response

# queries = generate_queries("王安石干了啥事，它有哪些作品")

@router.get("/generate_queries/{query}")
async def generate_queries_endpoint(query: str):
    return {"generated_queries": generate_queries(query)}


# Example usage:
# http://10.6.26.37:8002/generate_queries/王安石干了啥事，它有哪些作品